import csv
import os
from paddlenlp import Taskflow
import time
from colorama import Fore

from ocr import ocr_recognition

# 获取当前脚本文件的目录
script_directory = os.path.dirname(os.path.abspath(__file__))

# 构建 blank.jpg 的路径
blank_image_path = os.path.join(script_directory, "blank.jpg")

def ocr_nlp_processor(input_directory, output_directory, output_csv, prompt_to_header):
    """
    将OCR和NLP结合处理的函数。

    参数:
    - input_directory (str): 输入文件目录路径。
    - output_directory (str): 输出文件目录路径，用于保存识别结果文本文件。
    - output_csv (str): 输出CSV文件路径，用于保存识别结果的统计信息。
    - prompt_to_header (dict): 提示到表头的映射，指定需要识别的内容和对应的表头。

    返回:
    无，直接将识别结果写入到指定的输出文件中。
    """
    # 初始化NLP任务
    docPrompt = Taskflow('document_intelligence')
    
    # 获取表头列表，并添加文件名和识别时间列
    headers = ["文件名", "识别时间"] + [item["header"] for item in prompt_to_header.values()]
    
    # 如果CSV文件不存在，则创建文件并写入表头
    if not os.path.exists(output_csv):
        with open(output_csv, mode='w', newline='', encoding='utf-8') as file:
            writer = csv.writer(file)
            writer.writerow(headers)
    
    def ocr_callback(filename, result):
        """
        OCR识别的回调函数，用于处理OCR识别结果并执行NLP任务。
        
        参数:
        - filename (str): 被处理的文件名。
        - result (list): OCR识别结果，包括文本和位置等信息。
        """
        # 记录程序开始时间
        start_time = time.time()
        
        # 执行NLP任务
        nlpResult = docPrompt(
            [{"doc": blank_image_path, "prompt": list(prompt_to_header.keys()), "word_boxes": result}]
        )
        
        # 提取每个提示中概率最高的值，并进行个性化处理
        extracted_data = {header["header"]: "" for header in prompt_to_header.values()}
        for item in nlpResult:
            prompt = item['prompt']
            header_info = prompt_to_header[prompt]
            header = header_info["header"]
            process_value = header_info.get("process_value")
            if item['result']:
                best_result = max(item['result'], key=lambda x: x['prob'])
                value = best_result['value']
                if process_value:
                    value = process_value(value)
                extracted_data[header] = value
        
        # 计算程序执行时间
        end_time = time.time()
        total_time_seconds = end_time - start_time
        recognition_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(end_time))
        
        # 将文件名和识别时间添加到提取的数据中
        extracted_data["文件名"] = filename
        extracted_data["识别时间"] = recognition_time
        
        # 将提取的数据插入到CSV文件
        with open(output_csv, mode='a', newline='', encoding='utf-8') as file:
            writer = csv.writer(file)
            writer.writerow([extracted_data[header] for header in headers])
        
        # 打印文件处理完成的消息和运行总耗时
        print(Fore.GREEN + f"==> 文件 {filename} 的NLP识别完成，识别结果已写入CSV文件")
        print(Fore.GREEN + f"==> 本次NLP运行总耗时：{total_time_seconds:.2f} 秒")

    # 调用OCR识别函数处理输入目录中的文件
    ocr_recognition(input_directory, output_directory, ocr_callback)
